from QKDNetwork import QKDNetwork
from noise_provision import provision_with_noise_lp, provision_with_max_key_num
from lb_noise import calculateFlowPath
import random
from copy import deepcopy

# QKD请求的路由方案数据结构  [{"sd": [1, 2], "path": [1, 3, 2]}...]


class Compare:
    def __init__(self, net: QKDNetwork, avg_noise, threshold_noise):
        self.net: QKDNetwork = net
        self.avg_noise = avg_noise  # 平均噪声 0-19
        self.threshold_noise = threshold_noise  # 噪声阈值 0-19
        self.noised_qkd_networks = []
        self.generated_noise_networks()

    # 生成500个噪声网络
    def generated_noise_networks(self):
        for i in range(500):
            net = deepcopy(self.net)
            self.addNoise(net)
            self.noised_qkd_networks.append(net)

    # 向网络中注入噪声
    def inject_noise(self, net: QKDNetwork, idx):
        for node in net.G.nodes:
            net.G.nodes[node]["transmitter_rate"] = self.noised_qkd_networks[
                idx
            ].G.nodes[node]["transmitter_rate"]

    def getData(self):
        # return [self.rpsp_noise()]  # 不考虑噪声和考虑噪声
        return [self.rpsp(), self.rpsp_noise()]  # 不考虑噪声和考虑噪声

    def is_terminate_schedule(self, net: QKDNetwork):
        flag = True
        for sd_index, sd_item in enumerate(net.sd_list):
            # 还有没满足的SD需求，继续调度
            if sd_item[2] > sd_item[3]:
                flag = False
                break
        return flag

    def addNoise(self, net):
        # 向链路增加噪声
        # 将 net.G.nodes 的transmitter_rate 加上 diff_value
        for node in net.G.nodes:
            random_number = (random.random() - 0.5) / 50  # -0.1 ~ 0.1
            diff_value = random_number * self.avg_noise  # 噪声最大是 -2 ~ 2
            net.G.nodes[node]["transmitter_rate"] += diff_value
            if net.G.nodes[node]["transmitter_rate"] < 0:
                net.G.nodes[node]["transmitter_rate"] = 0.1

    def get_sd_idx_by_sd_pair(self, sd_start, sd_end):
        for i in range(len(self.net.sd_list)):
            if self.net.sd_list[i][0] == sd_start and self.net.sd_list[i][1] == sd_end:
                return i

    def rpsp(self, with_noise=False):
        terminate_schedule_flag = False
        max_time_slot = 0

        def sortSD(routeList):
            # 按照 需求量Demand - 已满足的密钥量satisfy_num 从小到大排序，然后按照路径长度从短到长排序
            routeList.sort(key=lambda x: (x["sd"][2] - x["sd"][3], len(x["path"])))

        tmp_net: QKDNetwork = deepcopy(self.noised_qkd_networks[0])

        route = None

        while not terminate_schedule_flag:
            max_time_slot += 1
            # 在一个时隙内，算路、算PS供应、在SD对上记账、在链路上记账
            if (route is None) or with_noise:
                route = calculateFlowPath(tmp_net)
            # 取出来一些路，发射密钥（一个时隙内）
            selected_direction = provision_with_max_key_num(route, tmp_net)
            for all_pop in selected_direction:
                pop_node_id = all_pop[0]
                dest_node_id = all_pop[1]
                transmitter_rate = tmp_net.G.nodes[pop_node_id]["transmitter_rate"]
                tmp_net.G[pop_node_id][dest_node_id][
                    "remain_key_num"
                ] += transmitter_rate
            # 先排序，然后满足请求的密钥
            sortSD(route)
            for i in route:
                sd = i["sd"]
                sd_idx = self.get_sd_idx_by_sd_pair(sd[0], sd[1])
                path = i["path"]

                max_remain_key_num = 1000000  # 该路径上剩余的最大密钥数量
                for j in range(len(path) - 1):
                    link_start, link_end = path[j], path[j + 1]
                    remain_key_num = tmp_net.G[link_start][link_end]["remain_key_num"]
                    max_remain_key_num = (
                        remain_key_num
                        if max_remain_key_num > remain_key_num
                        else max_remain_key_num
                    )
                need_key_num = i["sd"][2] - i["sd"][3]  # 实际需要再发射的密钥数量
                if need_key_num <= 0:
                    # 该源目节点对不需要密钥了
                    continue

                actual_emit_num = 0  # 实际发射的密钥数量

                if need_key_num >= max_remain_key_num:
                    actual_emit_num = max_remain_key_num

                if need_key_num < max_remain_key_num:
                    actual_emit_num = need_key_num

                # print("actual_emit_num: ", actual_emit_num)

                old_sd = tmp_net.sd_list[sd_idx]
                tmp_net.sd_list[sd_idx] = (
                    old_sd[0],
                    old_sd[1],
                    old_sd[2],
                    old_sd[3] + actual_emit_num,
                )
                i["sd"] = tmp_net.sd_list[sd_idx]

                for j in range(len(path) - 1):
                    link_start, link_end = path[j], path[j + 1]
                    tmp_net.G[link_start][link_end]["remain_key_num"] -= actual_emit_num
                    if tmp_net.G[link_start][link_end]["remain_key_num"] < 0:
                        print("error")

            total_remain_key_num = 0
            for edge in tmp_net.G.edges:
                total_remain_key_num += tmp_net.G[edge[0]][edge[1]]["remain_key_num"]

            print("remain key", total_remain_key_num)

            if max_time_slot > 50:
                print("shift")
                pass
            terminate_schedule_flag = self.is_terminate_schedule(tmp_net)
            self.inject_noise(tmp_net, max_time_slot)
        print(f"final time slot{'' if with_noise else ' optimal'}:", max_time_slot)
        return max_time_slot

    def rpsp_noise(self):
        self.rpsp(with_noise=True)
